import csv
import os
import re
import string

import mindspore
import mindspore.dataset as ds
import numpy as np


class AmazonData():
    """AMAZON数据集加载器

    加载AMAZON数据集并处理为一个Python迭代对象。

    """
    def __init__(self, path):
        self.path = path
        # 五个类别的评论分别放在五个列表里， 并以类别号作为索引
        self.review = [[] for _ in range(5)]
        # 统计五个类别各自样本数
        self.class_number = [0] * 5
        self._load()


    def _load(self):
        # 将数据加载至内存
        csvFile = open(self.path, "r")
        dict_reader = csv.DictReader(csvFile)
        trantab = str.maketrans(string.punctuation, ' '*len(string.punctuation))

        for row in dict_reader:
            review = row['review']
            label = int(np.float32(row['label']))
            label_onehot = [0] * 5
            label_onehot[label-1] = 1
            review = review.rstrip("\n\r").translate(trantab).lower()
            review = re.split(' |,', review.lower())
            # 将第 label-1 类的数据添加进对应位置
            self.review[label-1].append(review)
            # 这个类别的数量增加 1
            self.class_number[label-1] += 1
        # 统计最大类别个数
        self.max_class = max(self.class_number)

    def __getitem__(self, fake_idx):
        # 计算虚拟索引对应的标签， 标签 = 虚拟标签 除以 最大类别数 的商取整 比如一共250000个虚拟样本， 索引180000的
        # 标签为 180000 // 50000 = 3
        label = fake_idx // self.max_class
        # 计算 类内 的虚拟标签, 比如 虚拟标签 180000 对应第 3 类的第 180000 % 50000 = 30000 个 样本.
        fake_classidx = fake_idx % self.max_class 
        # 计算类内样本的真实标签， 真实标签 = 类内的虚拟标签 除以 类内真实样本数 的商取整
        # 比如第 3 类，第 30000 个样本对应的真实标签为 30000 % 第 3 类总样本数
        real_idx = fake_classidx % self.class_number[label]
        # 返回第 label 类，第 real_idx 个真实样本
        return self.review[label][real_idx], label

    def __len__(self):
        # 总样本数为最大类别数乘5
        return self.max_class * 5

def load_dataset(path):
    raw=AmazonData(path)
    raw=ds.GeneratorDataset(raw,column_names=["review", "label"], shuffle=False)
    return raw

# 加载词向量
def load_glove(glove_path,dimention):

    embeddings = []
    tokens = []
    with open(glove_path, encoding='utf-8') as gf:
        for glove in gf:
            word, embedding = glove.split(maxsplit=1)
            tokens.append(word)
            embeddings.append(np.fromstring(embedding, dtype=np.float32, sep=' '))
    # 添加 <unk>, <pad> 两个特殊占位符对应的embedding
    embeddings.append(np.random.rand(dimention))
    embeddings.append(np.zeros((dimention,), np.float32))

    vocab = ds.text.Vocab.from_list(tokens, special_tokens=["<unk>", "<pad>"], special_first=False)
    embeddings = np.array(embeddings).astype(np.float32)
    return vocab, embeddings

def data_preprocess(raw_data,vocab,embedding,batch_size):
    lookup_op = ds.text.Lookup(vocab, unknown_token='<unk>')
    pad_op = ds.transforms.c_transforms.PadEnd([400], pad_value=vocab.tokens_to_ids('<pad>'))
    type_cast_op = ds.transforms.c_transforms.TypeCast(mindspore.float32)
    data_train = raw_data.map(operations=[lookup_op, pad_op], input_columns=['review'])
    data_train = data_train.map(operations=[type_cast_op], input_columns=['label'])
    data_train = data_train.batch(batch_size, drop_remainder=True)
    
    return data_train